use strict;
use warnings;
use Scalar::Util qw(looks_like_number);
use POSIX;

require 'FilterData.pl';

open (IN_HD, "GENE_HD.txt");
open (IN_LD, "GENE_LD.txt");
open (OUT_TRAIN_HD, ' > GENE_TRAIN_HD.txt');
open (OUT_TEST_HD, ' > GENE_TEST_HD.txt');
open (OUT_TRAIN_LD, ' > GENE_TRAIN_LD.txt');
open (OUT_TEST_LD, ' > GENE_TEST_LD.txt');

my $nPos = 0;
my $nNeg = 0;
while(<IN_LD>)
{
    chomp;
    my $line_HD = $_;
    my @arr = split " ", $line_HD;
    my @t  = split ":", $arr[0];
    if($t[1] == 1)
    {
	$nPos++;
    }
    else
    {
	$nNeg++;
    }
}
close(IN_LD);
open (IN_LD, "GENE_LD.txt");

my $nPos_train = floor(3/4 * $nPos + 0.5);
my $nPos_test = $nPos - $nPos_train;

my $nNeg_train = floor(3/4 * $nNeg + 0.5);
my $nNeg_test = $nNeg - $nNeg_train;


my $iPos_train = 0;
my $iPos_test = 0;
my $iNeg_train = 0;
my $iNeg_test = 0;

while(<IN_LD>)
{
    chomp;
    my $line_LD = $_;
    my @arr = split " ", $line_LD;
    my @t  = split ":", $arr[0];
    my $line_HD = <IN_HD>;
        
    my $trainORtest = rand();
    
    if($trainORtest <= 0.5) #Throw in the train set
    {
	if($t[1] == 1) #If positive sample
	{
	    if($iPos_train < $nPos_train) #If train file is not full
	    {
		print OUT_TRAIN_LD "$line_LD\n";
		print OUT_TRAIN_HD "$line_HD";
		$iPos_train++;
	    }
	    else # Otherwise throw in the test set
	    {
		print OUT_TEST_LD "$line_LD\n";
		print OUT_TEST_HD "$line_HD";
		$iPos_test++;
	    }
	}
	else #If negative sample
	{
	    if($iNeg_train < $nNeg_train)
	    {
		print OUT_TRAIN_LD "$line_LD\n";
		print OUT_TRAIN_HD "$line_HD";
		$iNeg_train++;
	    }
	    else
	    {
		print OUT_TEST_LD "$line_LD\n";
		print OUT_TEST_HD "$line_HD";
		$iNeg_test++;
	    }
	}
    }
    else
    {
	if($t[1] == 1) #If positive sample
	{
	    if($iPos_test < $nPos_test) #If train file is not full
	    {
		print OUT_TEST_LD "$line_LD\n";
		print OUT_TEST_HD "$line_HD";
		$iPos_test++;
	    }
	    else # Otherwise throw in the test set
	    {
		print OUT_TRAIN_LD "$line_LD\n";
		print OUT_TRAIN_HD "$line_HD";
		$iPos_train++;
	    }
	}
	else #If negative sample
	{
	    if($iNeg_test < $nNeg_test)
	    {
		print OUT_TEST_LD "$line_LD\n";
		print OUT_TEST_HD "$line_HD";
		$iNeg_test++;
	    }
	    else
	    {
		print OUT_TRAIN_LD "$line_LD\n";
		print OUT_TRAIN_HD "$line_HD";
		$iNeg_train++;
	    }
	}
    }
}

close(IN_HD);
close(IN_LD);
close(OUT_TRAIN_HD);
close(OUT_TEST_HD);
close(OUT_TRAIN_LD);
close(OUT_TEST_LD);